#!/usr/bin/env python3
import rospy
import cv2
from cv_bridge import CvBridge
from sensor_msgs.msg import Image

class ImageStreamDisplay:
    def __init__(self):
        self.bridge = CvBridge()

        # 订阅彩色图像和深度图像话题
        rospy.Subscriber('/camera/color/image_raw', Image, self.color_callback)
        rospy.Subscriber('/camera/depth/image_raw', Image, self.depth_callback)

        # 初始化窗口
        cv2.namedWindow("Color Stream", cv2.WINDOW_NORMAL)
        cv2.namedWindow("Depth Stream", cv2.WINDOW_NORMAL)

        rospy.loginfo("ImageStreamDisplay node started, displaying color and depth streams.")

    def color_callback(self, msg):
        try:
            # 将ROS图像消息转换为OpenCV格式（BGR）
            cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
            # 显示彩色视频流
            cv2.imshow("Color Stream", cv_image)
            cv2.waitKey(1)  # 刷新窗口
        except Exception as e:
            rospy.logerr(f"Failed to process color image: {e}")

    def depth_callback(self, msg):
        try:
            # 将ROS深度图像消息转换为OpenCV格式
            depth_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough')
            # 归一化深度图像以便显示（将16位深度值映射到8位灰度图）
            depth_image_normalized = cv2.normalize(depth_image, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
            # 显示深度视频流
            cv2.imshow("Depth Stream", depth_image_normalized)
            cv2.waitKey(1)  # 刷新窗口
        except Exception as e:
            rospy.logerr(f"Failed to process depth image: {e}")

    def shutdown(self):
        # 节点关闭时销毁OpenCV窗口
        cv2.destroyAllWindows()
        rospy.loginfo("ImageStreamDisplay node shutdown.")

def main():
    rospy.init_node('image_stream_display', anonymous=True)
    display = ImageStreamDisplay()
    # 注册关闭钩子
    rospy.on_shutdown(display.shutdown)
    try:
        rospy.spin()
    except KeyboardInterrupt:
        rospy.loginfo("Shutting down image stream display node.")
        cv2.destroyAllWindows()

if __name__ == '__main__':
    main()